#!/usr/bin/env python3
import argparse
import os
import sys
import time
from datetime import datetime
from time import sleep

import grpc

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

SWITCH_TO_HOST_PORT = 1

# 创建统计目录
STATS_DIR = "tunnel_stats"
if not os.path.exists(STATS_DIR):
    os.makedirs(STATS_DIR)

# 定义隧道方向映射
TUNNEL_DIRECTIONS = {
    's1_to_s2': {'ingress_sw': 's1', 'egress_sw': 's2', 'tunnel_id': 100},
    's2_to_s1': {'ingress_sw': 's2', 'egress_sw': 's1', 'tunnel_id': 200},
    's1_to_s3': {'ingress_sw': 's1', 'egress_sw': 's3', 'tunnel_id': 101},
    's3_to_s1': {'ingress_sw': 's3', 'egress_sw': 's1', 'tunnel_id': 300},
    's2_to_s3': {'ingress_sw': 's2', 'egress_sw': 's3', 'tunnel_id': 201},
    's3_to_s2': {'ingress_sw': 's3', 'egress_sw': 's2', 'tunnel_id': 301},
}

def get_switch_to_switch_port(switch_name, tunnel_id):
    """
    根据交换机名称和隧道ID确定交换机之间的连接端口
    """
    # s1 的端口配置
    if switch_name == 's1':
        if tunnel_id == 100:  # s1 -> s2 (h1 -> h2)
            return 2
        elif tunnel_id == 101:  # s1 -> s3 (h1 -> h3)
            return 3
    # s2 的端口配置        
    elif switch_name == 's2':
        if tunnel_id == 200:  # s2 -> s1 (h2 -> h1)
            return 2
        elif tunnel_id == 201:  # s2 -> s3 (h2 -> h3)
            return 3
    # s3 的端口配置
    elif switch_name == 's3':
        if tunnel_id == 300:  # s3 -> s1 (h3 -> h1)
            return 2
        elif tunnel_id == 301:  # s3 -> s2 (h3 -> h2)
            return 3
    else:
        raise ValueError(f"Unknown switch or tunnel ID: {switch_name}, {tunnel_id}")

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr):
    """
    Installs three rules:
    1) An tunnel ingress rule on the ingress switch in the ipv4_lpm table that
       encapsulates traffic into a tunnel with the specified ID
    2) A transit rule on the ingress switch that forwards traffic based on
       the specified ID
    3) An tunnel egress rule on the egress switch that decapsulates traffic
       with the specified ID and sends it to the host
    """
    # 获取交换机间连接端口
    switch_to_switch_port = get_switch_to_switch_port(ingress_sw.name, tunnel_id)

    # 1) Tunnel Ingress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        })
    ingress_sw.WriteTableEntry(table_entry)
    print(f"Installed ingress tunnel rule on {ingress_sw.name} for tunnel {tunnel_id} to {dst_ip_addr}")

    # 2) Tunnel Transit Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_forward",
        action_params={
            "port": switch_to_switch_port
        })
    ingress_sw.WriteTableEntry(table_entry)
    print(f"Installed transit tunnel rule on {ingress_sw.name} for tunnel {tunnel_id} to port {switch_to_switch_port}")

    # 3) Tunnel Egress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": SWITCH_TO_HOST_PORT
        })
    egress_sw.WriteTableEntry(table_entry)
    print(f"Installed egress tunnel rule on {egress_sw.name} for tunnel {tunnel_id} from {dst_eth_addr}")

def read_counter_value(sw, counter_name, index):
    """
    读取特定计数器的值
    """
    try:
        for response in sw.ReadCounters(counter_name, index):
            for entity in response.entities:
                counter = entity.counter_entry
                return counter.data.packet_count, counter.data.byte_count
    except:
        pass
    return 0, 0

def print_tunnel_counters(p4info_helper, s1, s2, s3):
    """
    按照您要求的格式显示隧道计数器
    """
    print("\n" + "="*20 + " Reading tunnel counters " + "="*20)
    
    # 获取计数器ID
    ingress_counter_id = p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter")
    egress_counter_id = p4info_helper.get_counters_id("MyIngress.egressTunnelCounter")
    
    # 定义交换机映射
    switches = {'s1': s1, 's2': s2, 's3': s3}
    
    # 显示每个方向的隧道统计
    directions = [
        ('s1 -> s2', 100, 's1', 's2'),
        ('s2 -> s1', 200, 's2', 's1'),
        ('s1 -> s3', 101, 's1', 's3'),
        ('s3 -> s1', 300, 's3', 's1'),
        ('s2 -> s3', 201, 's2', 's3'),
        ('s3 -> s2', 301, 's3', 's2')
    ]
    
    for direction_name, tunnel_id, ingress_sw_name, egress_sw_name in directions:
        print(f"\n--- {direction_name} ---")
        
        # 读取入口计数器
        ingress_sw = switches[ingress_sw_name]
        ingress_packets, ingress_bytes = read_counter_value(ingress_sw, ingress_counter_id, tunnel_id)
        print(f"{ingress_sw_name} MyIngress.ingressTunnelCounter {tunnel_id}: {ingress_packets} packets ({ingress_bytes} bytes)")
        
        # 读取出口计数器
        egress_sw = switches[egress_sw_name]
        egress_packets, egress_bytes = read_counter_value(egress_sw, egress_counter_id, tunnel_id)
        print(f"{egress_sw_name} MyIngress.egressTunnelCounter {tunnel_id}: {egress_packets} packets ({egress_bytes} bytes)")
    
    print("\n" + "="*20 + " Finished " + "="*20)

def write_stats_to_files(p4info_helper, s1, s2, s3):
    """
    将统计信息写入文件
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # 获取计数器ID
    ingress_counter_id = p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter")
    egress_counter_id = p4info_helper.get_counters_id("MyIngress.egressTunnelCounter")
    
    # 定义交换机映射
    switches = {'s1': s1, 's2': s2, 's3': s3}
    
    # 为每个方向写入文件
    directions = [
        ('s1_to_s2', 100, 's1', 's2'),
        ('s2_to_s1', 200, 's2', 's1'),
        ('s1_to_s3', 101, 's1', 's3'),
        ('s3_to_s1', 300, 's3', 's1'),
        ('s2_to_s3', 201, 's2', 's3'),
        ('s3_to_s2', 301, 's3', 's2')
    ]
    
    for direction_name, tunnel_id, ingress_sw_name, egress_sw_name in directions:
        stats_file = os.path.join(STATS_DIR, f"{direction_name}.txt")
        
        # 读取入口和出口计数器
        ingress_sw = switches[ingress_sw_name]
        egress_sw = switches[egress_sw_name]
        
        ingress_packets, ingress_bytes = read_counter_value(ingress_sw, ingress_counter_id, tunnel_id)
        egress_packets, egress_bytes = read_counter_value(egress_sw, egress_counter_id, tunnel_id)
        
        # 写入文件
        with open(stats_file, 'a') as f:
            f.write(f"{timestamp}\n")
            f.write(f"{ingress_sw_name}发送{ingress_packets}个包，计数器编号为：{tunnel_id}\n")
            f.write(f"{egress_sw_name}收到{egress_packets}个包，计数器编号为：{tunnel_id}\n")
            f.write(f"{ingress_sw_name}->{egress_sw_name}\n\n")

def initialize_stats_files():
    """
    初始化统计文件，写入表头
    """
    directions = ['s1_to_s2', 's2_to_s1', 's1_to_s3', 's3_to_s1', 's2_to_s3', 's3_to_s2']
    
    for direction in directions:
        stats_file = os.path.join(STATS_DIR, f"{direction}.txt")
        with open(stats_file, 'w') as f:
            f.write(f"=== {direction.upper()} 隧道统计 ===\n")
            f.write("时间戳 | 发送包数 | 接收包数\n")
            f.write("=" * 40 + "\n")
    print("Statistics files initialized")

def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def printGrpcError(e):
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print("(%s)" % status_code.name, end=' ')
    traceback = sys.exc_info()[2]
    print("[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno))

def main(p4info_file_path, bmv2_file_path):
    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    # 初始化统计文件
    initialize_stats_files()

    try:
        # Create switch connections for s1, s2, and s3
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update message to establish this controller as master
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install the P4 program on all switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")

        print("\n=== Configuring Tunnel Rules ===")

        # 配置h1到h2的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s2, tunnel_id=100,
                         dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2")

        # 配置h1到h3的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s3, tunnel_id=101,
                         dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3")

        # 配置h2到h1的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s1, tunnel_id=200,
                         dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1")

        # 配置h2到h3的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s3, tunnel_id=201,
                         dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3")

        # 配置h3到h1的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s1, tunnel_id=300,
                         dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1")

        # 配置h3到h2的隧道规则
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s2, tunnel_id=301,
                         dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2")

        # 读取所有交换机的表规则
        print("\n=== Reading Table Rules ===")
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        print("\n=== Starting Counter Monitoring ===")
        print("Statistics will be saved to files in the 'tunnel_stats' directory")
        print("Press Ctrl+C to stop monitoring\n")
        
        cycle_count = 0
        while True:
            cycle_count += 1
            print(f"\n----- Cycle {cycle_count}: {datetime.now().strftime('%H:%M:%S')} -----")
            
            # 显示隧道计数器（按您要求的格式）
            print_tunnel_counters(p4info_helper, s1, s2, s3)
            
            # 写入统计文件
            write_stats_to_files(p4info_helper, s1, s2, s3)
            
            # 显示文件状态
            print(f"\nStatistics files updated in '{STATS_DIR}' directory:")
            for direction in ['s1_to_s2', 's2_to_s1', 's1_to_s3', 's3_to_s1', 's2_to_s3', 's3_to_s2']:
                file_path = os.path.join(STATS_DIR, f"{direction}.txt")
                if os.path.exists(file_path):
                    file_size = os.path.getsize(file_path)
                    print(f"  {direction}.txt: {file_size} bytes")
            
            sleep(5)  # 每5秒收集一次统计信息

    except KeyboardInterrupt:
        print("\n=== Final Statistics Summary ===")
        # 显示最终的统计信息
        print_tunnel_counters(p4info_helper, s1, s2, s3)
        print("\nShutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)